{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BN+conv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/media/pc/data/lxw/ai/tvm-book/doc/tutorials/frontend\n"
     ]
    }
   ],
   "source": [
    "%cd ../../..\n",
    "import set_env\n",
    "from d2py.utils.file import mkdir\n",
    "temp_dir = \".temp\"\n",
    "mkdir(temp_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Exported graph: graph(%data : Float(*, 3, 32, 32, strides=[3072, 1024, 32, 1], requires_grad=0, device=cpu),\n",
      "      %conv1.weight : Float(16, 3, 1, 1, strides=[3, 1, 1, 1], requires_grad=1, device=cpu),\n",
      "      %bn1.weight : Float(16, strides=[1], requires_grad=1, device=cpu),\n",
      "      %bn1.bias : Float(16, strides=[1], requires_grad=1, device=cpu),\n",
      "      %conv2.weight : Float(32, 16, 1, 1, strides=[16, 1, 1, 1], requires_grad=1, device=cpu)):\n",
      "  %bn1.running_var : Float(16, strides=[1], requires_grad=0, device=cpu) = onnx::Identity(%bn1.weight)\n",
      "  %bn1.running_mean : Float(16, strides=[1], requires_grad=0, device=cpu) = onnx::Identity(%bn1.bias)\n",
      "  %/conv1/Conv_output_0 : Float(*, 16, 32, 32, strides=[16384, 1024, 32, 1], requires_grad=0, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[1, 1], pads=[0, 0, 0, 0], strides=[1, 1], onnx_name=\"/conv1/Conv\"](%data, %conv1.weight), scope: __main__.Model::/torch.nn.modules.conv.Conv2d::conv1 # /media/pc/data/tmp/cache/conda/envs/xin/lib/python3.12/site-packages/torch/nn/modules/conv.py:456:0\n",
      "  %/act1/Relu_output_0 : Float(*, 16, 32, 32, strides=[16384, 1024, 32, 1], requires_grad=1, device=cpu) = onnx::Relu[onnx_name=\"/act1/Relu\"](%/conv1/Conv_output_0), scope: __main__.Model::/torch.nn.modules.activation.ReLU::act1 # /media/pc/data/tmp/cache/conda/envs/xin/lib/python3.12/site-packages/torch/nn/functional.py:1473:0\n",
      "  %/bn1/BatchNormalization_output_0 : Float(*, 16, 32, 32, strides=[16384, 1024, 32, 1], requires_grad=1, device=cpu) = onnx::BatchNormalization[epsilon=1.0000000000000001e-05, momentum=0.90000000000000002, training_mode=0, onnx_name=\"/bn1/BatchNormalization\"](%/act1/Relu_output_0, %bn1.weight, %bn1.bias, %bn1.running_mean, %bn1.running_var), scope: __main__.Model::/torch.nn.modules.batchnorm.BatchNorm2d::bn1 # /media/pc/data/tmp/cache/conda/envs/xin/lib/python3.12/site-packages/torch/nn/functional.py:2482:0\n",
      "  %output : Float(*, 32, 32, 32, strides=[32768, 1024, 32, 1], requires_grad=1, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[1, 1], pads=[0, 0, 0, 0], strides=[1, 1], onnx_name=\"/conv2/Conv\"](%/bn1/BatchNormalization_output_0, %conv2.weight), scope: __main__.Model::/torch.nn.modules.conv.Conv2d::conv2 # /media/pc/data/tmp/cache/conda/envs/xin/lib/python3.12/site-packages/torch/nn/modules/conv.py:456:0\n",
      "  return (%output)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "\n",
    "class Model(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv2d(3, 16, 1, 1, 0, bias=False, groups=1)\n",
    "        self.act1  = nn.ReLU()\n",
    "        self.bn1   = nn.BatchNorm2d(num_features=16)\n",
    "        self.conv2 = nn.Conv2d(16, 32, 1, 1, 0, bias=False, groups=1)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.act1(x)\n",
    "        x = self.bn1(x)\n",
    "        x = self.conv2(x)\n",
    "        return x\n",
    "\n",
    "shape = 1, 3, 32, 32\n",
    "x = torch.rand(*shape)\n",
    "\n",
    "torch_model = Model()\n",
    "# 导出模型\n",
    "output_name = \"bn-conv\"\n",
    "torch.onnx.export(\n",
    "    torch_model,               # torch 模型\n",
    "    x,                         # 模型输入或者对于多个输入，使用元组\n",
    "    f\"{temp_dir}/{output_name}.onnx\",               # 模型保存的位置（可以是文件或类似文件的对象）\n",
    "    export_params=True,        # 将训练后的参数权重存储在模型文件内\n",
    "    opset_version=17,          # 导出模型的 ONNX 版本\n",
    "    do_constant_folding=True,  # 是否执行常量折叠以进行优化\n",
    "    input_names = ['data'],    # 模型的输入名称\n",
    "    output_names = ['output'], # 模型的输出名称\n",
    "    # keep_initializers_as_inputs=True,\n",
    "    # export_modules_as_functions=True,\n",
    "    verbose=True,\n",
    "    dynamic_axes={'data' : {0 : 'batch_size'},    # 可变长度的轴\n",
    "                  'output' : {0 : 'batch_size'}}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](images/bn-conv.jpg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float32] span<span style=\"color: #AA22FF; font-weight: bold\">=/</span>conv1<span style=\"color: #AA22FF; font-weight: bold\">/</span>Conv<span style=\"color: #AA22FF; font-weight: bold\">.</span>data:<span style=\"color: #008000\">0</span>:<span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data, meta[relay<span style=\"color: #AA22FF; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">0</span>] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32] span<span style=\"color: #AA22FF; font-weight: bold\">=/</span>conv1<span style=\"color: #AA22FF; font-weight: bold\">/</span>Conv<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv1<span style=\"color: #AA22FF; font-weight: bold\">.</span>weight:<span style=\"color: #008000\">0</span>:<span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], channels<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">16</span>, kernel_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>]) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float32] span<span style=\"color: #AA22FF; font-weight: bold\">=/</span>conv1<span style=\"color: #AA22FF; font-weight: bold\">/</span>Conv:<span style=\"color: #008000\">0</span>:<span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float32] span<span style=\"color: #AA22FF; font-weight: bold\">=/</span>act1<span style=\"color: #AA22FF; font-weight: bold\">/</span>Relu:<span style=\"color: #008000\">0</span>:<span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>batch_norm(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, meta[relay<span style=\"color: #AA22FF; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">1</span>] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] span<span style=\"color: #AA22FF; font-weight: bold\">=</span>Identity_0<span style=\"color: #AA22FF; font-weight: bold\">.</span>bn1<span style=\"color: #AA22FF; font-weight: bold\">.</span>weight:<span style=\"color: #008000\">0</span>:<span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, meta[relay<span style=\"color: #AA22FF; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">2</span>] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] span<span style=\"color: #AA22FF; font-weight: bold\">=</span>Identity_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>bn1<span style=\"color: #AA22FF; font-weight: bold\">.</span>bias:<span style=\"color: #008000\">0</span>:<span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, meta[relay<span style=\"color: #AA22FF; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">3</span>] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] span<span style=\"color: #AA22FF; font-weight: bold\">=</span>Identity_1:<span style=\"color: #008000\">0</span>:<span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, meta[relay<span style=\"color: #AA22FF; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">4</span>] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>), float32] span<span style=\"color: #AA22FF; font-weight: bold\">=</span>Identity_0:<span style=\"color: #008000\">0</span>:<span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32], Tensor[(<span style=\"color: #008000\">16</span>), float32]) span<span style=\"color: #AA22FF; font-weight: bold\">=/</span>bn1<span style=\"color: #AA22FF; font-weight: bold\">/</span>BatchNormalization:<span style=\"color: #008000\">0</span>:<span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2.0</span> <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float32] span<span style=\"color: #AA22FF; font-weight: bold\">=/</span>bn1<span style=\"color: #AA22FF; font-weight: bold\">/</span>BatchNormalization:<span style=\"color: #008000\">0</span>:<span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, meta[relay<span style=\"color: #AA22FF; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">5</span>] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), float32] span<span style=\"color: #AA22FF; font-weight: bold\">=/</span>conv2<span style=\"color: #AA22FF; font-weight: bold\">/</span>Conv<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2<span style=\"color: #AA22FF; font-weight: bold\">.</span>weight:<span style=\"color: #008000\">0</span>:<span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], channels<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">32</span>, kernel_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>]) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float32] span<span style=\"color: #AA22FF; font-weight: bold\">=/</span>conv2<span style=\"color: #AA22FF; font-weight: bold\">/</span>Conv:<span style=\"color: #008000\">0</span>:<span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import onnx\n",
    "import tvm\n",
    "from tvm import relay\n",
    "onnx_model = onnx.load(f\"{temp_dir}/{output_name}.onnx\")\n",
    "mod, params = relay.frontend.from_onnx(onnx_model, {\"data\": shape}, freeze_params=True)\n",
    "# with tvm.transform.PassContext(opt_level=3):\n",
    "#     mod = relay.quantize.prerequisite_optimize(mod, params)\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float32] span<span style=\"color: #AA22FF; font-weight: bold\">=/</span>conv1<span style=\"color: #AA22FF; font-weight: bold\">/</span>Conv<span style=\"color: #AA22FF; font-weight: bold\">.</span>data:<span style=\"color: #008000\">0</span>:<span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span>data, <span style=\"color: #008000\">16</span>f <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>float32 <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> round(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> clip(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>, a_min<span style=\"color: #AA22FF; font-weight: bold\">=-</span><span style=\"color: #008000\">127</span>f, a_max<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">127</span>f) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> cast(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>, meta[relay<span style=\"color: #AA22FF; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">0</span>] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], channels<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">16</span>, kernel_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], out_dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>relu(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">6</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fixed_point_multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">5</span>, multiplier<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">0</span>, shift<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">0</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> cast(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">6</span>, dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">8</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> add(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">7</span>, meta[relay<span style=\"color: #AA22FF; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">1</span>] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">9</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> cast(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">8</span>, dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int64&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), int64] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">10</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fixed_point_multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">9</span>, multiplier<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">0</span>, shift<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">0</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), int64] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">11</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> clip(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">10</span>, a_min<span style=\"color: #AA22FF; font-weight: bold\">=-</span><span style=\"color: #008000\">127</span>f, a_max<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">127</span>f) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), int64] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">12</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> cast(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">11</span>, dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">13</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> cast(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">12</span>, dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">14</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> annotation<span style=\"color: #AA22FF; font-weight: bold\">.</span>stop_fusion(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">13</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">15</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>conv2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">14</span>, meta[relay<span style=\"color: #AA22FF; font-weight: bold\">.</span>Constant][<span style=\"color: #008000\">2</span>] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">16</span>, <span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], channels<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">32</span>, kernel_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], out_dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">16</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> cast(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">15</span>, dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int64&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), int64] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">17</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fixed_point_multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">16</span>, multiplier<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">2146285056</span>, shift<span style=\"color: #AA22FF; font-weight: bold\">=-</span><span style=\"color: #008000\">9</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), int64] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">18</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> clip(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">17</span>, a_min<span style=\"color: #AA22FF; font-weight: bold\">=-</span><span style=\"color: #008000\">127</span>f, a_max<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">127</span>f) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), int64] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">19</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> cast(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">18</span>, dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), int32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">20</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> cast(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">19</span>, dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;int8&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">21</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> annotation<span style=\"color: #AA22FF; font-weight: bold\">.</span>stop_fusion(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">20</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">22</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> cast(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">21</span>, dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  multiply(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">22</span>, <span style=\"color: #008000\">0.0625</span>f <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>float32 <span style=\"color: #AA22FF; font-weight: bold\">*/</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>), float32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "with tvm.transform.PassContext(opt_level=3):\n",
    "    with relay.quantize.qconfig(\n",
    "        skip_conv_layers=[],\n",
    "        # calibrate_mode=\"kl_divergence\", \n",
    "        weight_scale=\"max\",\n",
    "        # round_for_shift=True,\n",
    "        # rounding=\"TONEAREST\", # \"UPWARD\" or \"TONEAREST\"\n",
    "        # calibrate_skip_layers=[],\n",
    "        skip_dense_layer=False,\n",
    "    ):\n",
    "        qmod = relay.quantize.quantize(mod, params)\n",
    "qmod.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "xin",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
